import numpy as np
from train_config import get_candidates, divide_candidates, get_statistical_function_and_logging_str, get_info
from easy_io import read_pkl_file, write_pkl_file
from itertools import chain
from find_best_epoch import find_best_epoch


def predict(candidates, probs):
    probs = np.asarray(probs)
    assert probs.shape == (len(candidates), 2)
    for c, p in zip(candidates, probs):
        c.update(prob=float(p[1]))
    return candidates


def main(candidate_file, result_files, valid_folds_lst, info_file, outfile):
    candidates = get_candidates(candidate_file)
    candidate_lst = []
    for valid_folds, result_file in zip(valid_folds_lst, result_files):
        valid_probs = read_pkl_file(result_file)
        _, valid_candidates = divide_candidates(candidates, valid_folds)
        candidate_lst.append(predict(valid_candidates, valid_probs))
        print(get_statistical_function_and_logging_str(valid_candidates, *get_info(info_file, valid_folds))[0](valid_probs)['recalls'])
    candidates = list(chain.from_iterable(candidate_lst))
    write_pkl_file(outfile, candidates)


if __name__ == '__main__':
    # main(
    #     candidate_file='lidc_kaggle_candidates_v5.pkl',
    #     result_files=[
    #         '/data_4t/huzq/kaggle_results/cddv5_revised/epoch-0116.pkl',
    #         '/data_4t/huzq/kaggle_results/cddv5_revised_valid1_v5/epoch-0088.pkl',
    #         '/data_4t/huzq/kaggle_results/cddv5_revised_valid2_v2/epoch-0138.pkl',
    #         '/data_4t/huzq/kaggle_results/cddv5_revised_valid3/epoch-0140.pkl',
    #         '/data_4t/huzq/kaggle_results/cddv5_revised_valid4/epoch-0130.pkl',
    #     ],
    #     valid_folds_lst=[
    #         [0],
    #         [1],
    #         [2],
    #         [3],
    #         [4],
    #     ],
    #     info_file='lidc_kaggle_info_v1.pkl',
    #     outfile='/data_4t/Kaggle/candidates/3dcnn_candidates_list_v5.pkl',
    # )

    # epoches = [140, 252, 270, 62, 139]
    # main(
    #     candidate_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v9.pkl',
    #     result_files=[
    #         '/data_4t/huzq/kaggle_results/cddv9_valid{i}/epoch-{epoch:04d}.pkl'.format(i=i, epoch=epoches[i])
    #         for i in range(len(epoches))
    #     ],
    #     valid_folds_lst=[
    #         [0],
    #         [1],
    #         [2],
    #         [3],
    #         [4],
    #     ],
    #     info_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_info_v2.pkl',
    #     outfile='/data_4t/Kaggle/candidates/3dcnn_candidates_list_v9.pkl',
    # )

    # result_files = [
    #         '/data_4t/huzq/kaggle_results/cddv9_vgg13shortcutsv2_bs16_valid{i}/epoch-{epoch:04d}.pkl'.format(i=i, epoch=int(find_best_epoch('/ssd_1t/huzq/kaggle_log/cddv9_vgg13shortcutsv2_bs16_valid{i}.log'.format(i=i))))
    #         for i in range(5)
    # ]
    # print(result_files)
    # main(
    #     candidate_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v9.pkl',
    #     result_files=result_files,
    #     valid_folds_lst=[
    #         [0],
    #         [1],
    #         [2],
    #         [3],
    #         [4],
    #     ],
    #     info_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_info_v2.pkl',
    #     outfile='/data_4t/Kaggle/candidates/3dcnn_candidates_list_v9_vgg13shortcutsv2_bs16.pkl',
    # )

    # result_files = [
    #     '/data_4t/huzq/kaggle_results/cddv9_vgg10_valid{i}/epoch-{epoch:04d}.pkl'.format(i=i, epoch=int(
    #         find_best_epoch('/ssd_1t/huzq/kaggle_log/cddv9_vgg10_valid{i}.log'.format(i=i))))
    #     for i in range(5)
    #     ]
    # print(result_files)
    # main(
    #     candidate_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v9.pkl',
    #     result_files=result_files,
    #     valid_folds_lst=[
    #         [0],
    #         [1],
    #         [2],
    #         [3],
    #         [4],
    #     ],
    #     info_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_info_v2.pkl',
    #     outfile='/data_4t/Kaggle/candidates/3dcnn_candidates_list_v9_vgg10.pkl',
    # )

    result_files = [
        '/data_4t/huzq/kaggle_results/cddv15_v2_valid{i}/epoch-{epoch:04d}.pkl'.format(i=i, epoch=int(
            find_best_epoch('/ssd_1t/huzq/kaggle_log/cddv15_v2_valid{i}.log'.format(i=i))))
        for i in range(4)
        ]
    print(result_files)
    main(
        candidate_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_candidates_v15.pkl',
        result_files=result_files,
        valid_folds_lst=[
            [0],
            [1],
            [2],
            [3],
        ],
        info_file='/ssd_1t/huzq/kaggle_data/lidc_kaggle_info_v4.pkl',
        outfile='/data_4t/Kaggle/candidates/3dcnn_cddv15_v2_lidc_kaggle_candidates_v15.pkl',
    )
